import math
from sys import maxsize
import matplotlib.pyplot as plt
from matplotlib import animation
from matplotlib.colors import ListedColormap
import numpy as np

# 表示地图上每个格子的状态
class State(object):
    def __init__(self, x, y):
        self.x = x  # 行坐标
        self.y = y  # 列坐标
        self.parent = None  # 回溯路径的前驱节点
        self.state = "."  # 状态符号：. 空地，# 障碍，s 起点，e 终点，* 路径
        self.t = "new"  # 标记状态：new/open/close
        self.h = 0  # 启发式代价值
        self.k = 0  # 最小路径估价函数

    def cost(self, state):
        # 如果有障碍返回极大值，否则计算欧几里得距离
        if self.state == "#" or state.state == "#":
            return maxsize
        return math.hypot(self.x - state.x, self.y - state.y)

    def set_state(self, state):
        # 设置格子状态
        if state not in ["s", ".", "#", "e", "*"]:
            return
        self.state = state


# 表示整个地图，包括网格和障碍
class Map(object):
    def __init__(self, row, col):
        self.row = row
        self.col = col
        self.map = self.init_map()

    def init_map(self):
        # 初始化地图为 State 对象的二维列表
        return [[State(i, j) for j in range(self.col)] for i in range(self.row)]

    def get_neighbers(self, state):
        # 获取一个格子周围8个方向的邻居（包括对角）
        state_list = []
        for i in [-1, 0, 1]:
            for j in [-1, 0, 1]:
                if i == 0 and j == 0:
                    continue
                nx, ny = state.x + i, state.y + j
                if 0 <= nx < self.row and 0 <= ny < self.col:
                    state_list.append(self.map[nx][ny])
        return state_list

    def set_obstacle(self, point_list):
        # 设置地图上的障碍
        for x, y in point_list:
            if 0 <= x < self.row and 0 <= y < self.col:
                self.map[x][y].set_state("#")


# 实现 D* 路径规划算法
class Dstar(object):
    def __init__(self, maps):
        self.map = maps
        self.open_list = set()  # 存放待处理节点
        self.frames = []  # 存储每帧动画状态
        self.start = None
        self.end = None

    def process_state(self):
        # 核心状态处理函数，根据 D* 三种情况更新路径信息
        x = self.min_state()  # 获取 open list 中代价最小的节点
        if x is None:
            return -1  # 如果没有节点，则返回 -1
        k_old = self.get_kmin()  # 获取当前代价最小的 k 值
        self.remove(x)  # 将节点 x 从 open list 移除
        if k_old < x.h:  # 如果当前 k 小于 x 的启发式代价
            # 处理邻居节点，更新路径和代价
            for y in self.map.get_neighbers(x):
                if y.h <= k_old and x.h > y.h + x.cost(y):
                    x.parent = y
                    x.h = y.h + x.cost(y)
        elif k_old == x.h:  # 如果 k 值相等，更新节点
            for y in self.map.get_neighbers(x):
                if y.t == "new" or (y.parent == x and y.h != x.h + x.cost(y)) or (y.parent != x and y.h > x.h + x.cost(y)):
                    y.parent = x
                    self.insert(y, x.h + x.cost(y))
        else:  # 如果 k 值更大，更新邻居节点
            for y in self.map.get_neighbers(x):
                if y.t == "new" or (y.parent == x and y.h != x.h + x.cost(y)):
                    y.parent = x
                    self.insert(y, x.h + x.cost(y))
                elif y.parent != x and y.h > x.h + x.cost(y):
                    self.insert(y, x.h)
                elif y.parent != x and x.h > y.h + x.cost(y) and y.t == "close" and y.h > k_old:
                    self.insert(y, y.h)
        return self.get_kmin()  # 返回最小的 k 值

    def min_state(self):
        # 获取 open list 中 k 最小的节点
        return min(self.open_list, key=lambda x: x.k) if self.open_list else None

    def get_kmin(self):
        # 获取 open list 中最小的 k 值
        return min((x.k for x in self.open_list), default=-1)

    def insert(self, state, h_new):
        # 将节点插入 open list，并设置其状态
        if state.t == "new":
            state.k = h_new  # 如果是新节点，则直接设置 k 值
        elif state.t == "open":
            state.k = min(state.k, h_new)  # 如果节点已经在 open list 中，取较小的 k 值
        elif state.t == "close":
            state.k = min(state.h, h_new)  # 如果节点已经关闭，则取较小的代价
        state.h = h_new
        state.t = "open"  # 设置为 open 状态
        self.open_list.add(state)  # 将节点添加到 open list

    def remove(self, state):
        # 将节点从 open list 移除
        if state.t == "open":
            state.t = "close"  # 设置为 close 状态
        self.open_list.remove(state)

    def modify_cost(self, x):
        # 修改节点的代价，触发重新插入 open list
        if x.t == "close":
            self.insert(x, x.parent.h + x.cost(x.parent))  # 重新插入并更新代价

    def run(self, start, end):
        path_length = 0  # 初始化路径长度统计
        path_cost = 0    # 初始化路径移动代价
        self.start = start
        self.end = end
        self.open_list.add(end)  # 从终点开始向起点反推
        while True:
            self.process_state()  # 处理当前状态
            if start.t == "close":  # 如果起点已关闭，说明找到路径
                break
        start.set_state("s")  # 设置起点状态
        s = start
        # 回溯路径并设置状态
        while s != end:
            s.set_state("s")
            path_length += 1  # 每步路径加一
            path_cost += s.cost(s.parent) if s.parent else 0  # 累加代价
            self.capture_frame()  # 捕捉当前帧
            s = s.parent
        s.set_state("e")  # 设置终点状态
        self.capture_frame()

        # 模拟新增障碍物触发重规划
        self.map.set_obstacle([(9, i) for i in range(3, 9)])
        self.capture_frame()

        print(f"初始路径长度：{path_length}")
        print(f"总移动代价：{path_cost:.2f}")

        tmp = start
        while tmp != end:
            tmp.set_state("*")  # 标记路径
            self.capture_frame()
            if tmp.parent.state == "#":  # 如果路径节点为障碍，重新规划
                self.modify(tmp)
                continue
            tmp = tmp.parent
        tmp.set_state("e")  # 设置终点状态
        self.capture_frame()

    def modify(self, state):
        # 重新修改路径代价并重规划
        self.modify_cost(state)
        while self.process_state() < state.h:  # 直到代价调整完毕
            pass

    def capture_frame(self):
        # 保存当前地图状态作为动画帧
        self.frames.append(self.get_frame_array())

    def get_frame_array(self):
        # 将状态字符转为颜色编码矩阵
        color_map = {'.': 0, '#': 1, 's': 2, 'e': 3, '*': 4}
        array = [[color_map.get(self.map.map[i][j].state, 0)
                  for j in range(self.map.col)] for i in range(self.map.row)]
        # 强调标记起点终点
        if self.start:
            array[self.start.x][self.start.y] = 2
        if self.end:
            array[self.end.x][self.end.y] = 3
        return array


# 绘制动画并保存为 gif
def animate_map(frames, save_path="dstar_path.gif"):
    fig, ax = plt.subplots()
    fig.subplots_adjust(left=0, right=1, bottom=0, top=1)  # 去除边距
    cmap = ListedColormap([
        '#add8e6',  # 浅蓝色 - 背景
        '#ff0000',  # 红色 - 障碍物
        '#0000ff',  # 蓝色 - 起点
        '#ff1493',  # 红色 - 终点（洋红）
        '#32cd32'   # 绿色 - 路径
    ])
    im = ax.imshow(np.array(frames[0]), cmap=cmap, vmin=0, vmax=4)
    ax.axis('off')

    def update(i):
        im.set_array(np.array(frames[i]))
        return [im]

    ani = animation.FuncAnimation(fig, update, frames=len(frames), interval=200, blit=True)
    ani.save(save_path, writer='pillow', fps=5)
    print(f"动画已保存为 GIF：{save_path}")
    plt.close(fig)


# 主函数：构建地图、设置障碍、运行 D*、生成动画
if __name__ == '__main__':
    m = Map(20, 20)
    m.set_obstacle([(4, 3), (4, 4), (4, 5), (4, 6), (5, 3), (6, 3), (7, 3)])  # 设置障碍
    start = m.map[1][2]  # 设置起点
    end = m.map[17][11]  # 设置终点
    dstar = Dstar(m)
    dstar.run(start, end)  # 路径规划
    animate_map(dstar.frames, save_path="dstar_path.gif")  # 动画保存
